import os
import json
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI

load_dotenv()
llm = ChatOpenAI(
    model_name="qwen-plus",
    api_key=os.getenv("DASHSCOPE_API_KEY"),
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
    temperature=0  # 0表示输出更确定，适合需要精确决策的场景
)

# 2. 定义结构化工具（天气+酒店）
from langchain_core.tools import tool


# 工具1：天气查询（多参数格式）
@tool(
    description="查询气温，需要提供城市和日期参数"
)
def WeatherTool(city: str, date: str) -> str:
    """
    查询指定城市和日期的天气情况

    Args:
        city: 城市名称
        date: 日期，格式为YYYY-MM-DD

    Returns:
        天气信息字符串
    """
    try:
        city = city.strip()
        date = date.strip()
        weather_data = {("北京", "2024-10-15"): "23℃", ("上海", "2024-10-15"): "26℃"}
        return f"{city} {date}：{weather_data.get((city, date), '暂无数据')}"
    except Exception as e:
        return f"天气查询出错：{str(e)}"


# 工具2：酒店预订（多参数格式）
@tool(
    description="预订酒店，需要提供城市、日期、房间类型和客人数量参数"
)
def HotelTool(city: str, date: str, room_type: str, guests: int) -> str:
    """
    预订酒店房间

    Args:
        city: 城市名称
        date: 入住日期，格式为YYYY-MM-DD
        room_type: 房间类型（如：大床房、双床房等）
        guests: 客人数量

    Returns:
        预订结果信息
    """
    try:
        # 参数校验
        if not all([city, date, room_type]):
            return "缺少必填参数！需要提供城市、日期和房间类型"

        if guests <= 0:
            return "客人数量必须大于0"

        # 模拟预订
        order_id = hash(f"{city}{date}{room_type}{guests}") % 1000000
        return f"已预订 {city} {date} {room_type}（{guests}人），订单号：H{order_id}"
    except Exception as e:
        return f"酒店预订出错：{str(e)}"


tools = [WeatherTool, HotelTool]

# 3. 初始化结构化记忆（扩展参数存储）
from langchain.memory import ConversationBufferMemory

memory = ConversationBufferMemory(
    memory_key="chat_history",
    return_messages=True,
    output_key="output"
)

# 4. 定义结构化 Prompt（支持跨工具参数复用）
from langchain_core.prompts import PromptTemplate

template = """
你是结构化助手，支持天气查询和酒店预订。

### 重要规则：
1. 仔细阅读历史对话，寻找城市和日期信息
2. 当用户要求订酒店时，如果历史对话中有城市和日期，必须使用这些信息
3. 不要要求用户重新提供已经在历史对话中出现的信息

### 工具使用：
- WeatherTool(city, date): 查询天气
- HotelTool(city, date, room_type, guests): 预订酒店

### 参数提取方法：
1. 从历史对话中查找城市名称（如：北京、上海等）
2. 从历史对话中查找日期（如：2024-10-15等）
3. 从当前用户输入中提取房间类型和客人数量
4. 如果历史对话中有城市和日期，直接使用；如果没有，才询问用户

### 具体示例：
历史对话：用户问"查询北京2024-10-15的天气"
当前输入：用户说"帮我订酒店，大床房，2人"
正确做法：使用city="北京", date="2024-10-15", room_type="大床房", guests=2

历史对话：
{chat_history}

当前用户输入：{input}

可用工具：
{tools}

请仔细分析历史对话，提取城市和日期信息，然后调用相应工具。
"""

prompt = PromptTemplate(
    template=template,
    input_variables=["input", "chat_history", "tools"],
    partial_variables={
        "tools": "\n".join([f"- {tool.name}：{tool.description}" for tool in tools]),
    }
)

# 5. 创建 Agent 与 Executor
from langchain.agents import initialize_agent, AgentType, AgentExecutor, StructuredChatAgent

agent = StructuredChatAgent.from_llm_and_tools(
    llm=llm,
    tools=tools,
    prompt=prompt,
    verbose=True,
    memory=memory
)

agent_executor = AgentExecutor(
    agent=agent,
    tools=tools,
    verbose=True,
    handle_parsing_errors="参数格式错误！请按工具示例重新提供",
    max_iterations=4,
    memory=memory
)


# 6. 自定义参数提取函数
def extract_parameters_from_history(chat_history, current_input):
    """从历史对话中提取城市和日期参数"""
    city = None
    date = None

    # 从历史对话中查找城市和日期
    for message in chat_history:
        content = message.content if hasattr(message, 'content') else str(message)
        # 查找城市
        if "北京" in content:
            city = "北京"
        elif "上海" in content:
            city = "上海"
        # 查找日期
        if "2024-10-15" in content:
            date = "2024-10-15"

    return city, date


# 7. 测试多轮对话与参数复用
if __name__ == "__main__":
    # 测试1：查天气（存储参数）
    print("===== 测试1：查天气 =====")
    print("Agent记忆1：", memory.buffer)
    res1 = agent_executor.invoke({"input": "查询北京2024-10-15的天气"})
    print("Agent回答：", res1["output"], "\n")

    # 测试2：订酒店（复用天气参数：北京,2024-10-15）
    print("===== 测试2：订酒店（复用参数） =====")
    print("Agent记忆2：", memory.buffer)

    # 手动提取参数并构建输入
    city, date = extract_parameters_from_history(memory.buffer, "帮我订酒店，大床房，2人")
    if city and date:
        enhanced_input = f"帮我订酒店，大床房，2人。城市：{city}，日期：{date}"
        print(f"提取的参数：城市={city}, 日期={date}")
        print(f"增强输入：{enhanced_input}")
        res2 = agent_executor.invoke({"input": enhanced_input})
    else:
        res2 = agent_executor.invoke({"input": "帮我订酒店，大床房，2人"})
    print("Agent回答：", res2["output"], "\n")

    # 测试3：查新城市天气（更新参数）
    print("===== 测试3：查新城市天气 =====")
    print("Agent记忆3：", memory.buffer)
    res3 = agent_executor.invoke({"input": "查询上海2024-10-15的天气"})
    print("Agent回答：", res3["output"], "\n")

    # 测试4：再次订酒店（复用新参数：上海,2024-10-15）
    print("===== 测试4：再次订酒店（复用新参数） =====")
    print("Agent记忆4：", memory.buffer)

    # 手动提取参数并构建输入
    city, date = extract_parameters_from_history(memory.buffer, "帮我订酒店，双床房，3人")
    if city and date:
        enhanced_input = f"帮我订酒店，双床房，3人。城市：{city}，日期：{date}"
        print(f"提取的参数：城市={city}, 日期={date}")
        print(f"增强输入：{enhanced_input}")
        res4 = agent_executor.invoke({"input": enhanced_input})
    else:
        res4 = agent_executor.invoke({"input": "帮我订酒店，双床房，3人"})
    print("Agent回答：", res4["output"])